#include <cstdio>
#include <memory>
#include <opencv2/opencv.hpp>
#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/model.h"
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/model_builder.h"
#include "tensorflow/lite/schema/schema_generated.h"

// 辅助函数：检查是否为车道线像素
bool is_lane_pixel(const float* instance_seg_data, int idx, int num_channels) {
    float sum = 0.0f;
    for (int c = 0; c < num_channels; ++c) {
        sum += instance_seg_data[idx * num_channels + c];
    }
    return (sum / num_channels) > -0.2f;
}

// 确定转向方向
std::string determine_turn_direction(const float* instance_seg_data, 
                                   int height, int width,
                                   int& left_count,
                                   int& right_count) {
    left_count = 0;
    right_count = 0;
    const int center_x = width / 2;
    const int num_channels = 4;
    
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            if (is_lane_pixel(instance_seg_data, y * width + x, num_channels)) {
                if (x < center_x) {
                    left_count++;
                } else {
                    right_count++;
                }
            }
        }
    }
    
    float total_pixels = left_count + right_count;
    if (total_pixels < 100) {
        return "Keep Straight";
    }
    
    float difference_percent = abs(left_count - right_count) / total_pixels;
    const float TURN_THRESHOLD = 0.15;
    
    if (difference_percent < TURN_THRESHOLD) {
        return "Keep Straight";
    } else if (left_count > right_count) {
        return "Turn Right";
    } else {
        return "Turn Left";
    }
}

// 图像预处理
cv::Mat preprocess_image(const cv::Mat& input_image) {
    cv::Mat resized_image;
    cv::resize(input_image, resized_image, cv::Size(512, 256), 0, 0, cv::INTER_LINEAR);
    cv::Mat float_img;
    resized_image.convertTo(float_img, CV_32F, 1.0/127.5, -1.0);
    return float_img;
}

// 后处理输出
void postprocess_outputs(const float* instance_seg_data,
                        int height, int width,
                        const cv::Mat& original_image) {
    if (!instance_seg_data) {
        throw std::runtime_error("Null instance seg data received");
    }

    // 创建实例掩码
    cv::Mat instance_mask(height, width, CV_8UC1, cv::Scalar(0));
    const int num_channels = 4;

    // 找到每个通道的最大最小值
    std::vector<float> min_vals(num_channels, std::numeric_limits<float>::max());
    std::vector<float> max_vals(num_channels, std::numeric_limits<float>::lowest());
    
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            for (int c = 0; c < num_channels; ++c) {
                float val = instance_seg_data[(y * width + x) * num_channels + c];
                min_vals[c] = std::min(min_vals[c], val);
                max_vals[c] = std::max(max_vals[c], val);
            }
        }
    }

    // 合并所有通道的信息并进行归一化
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            float normalized_sum = 0;
            for (int c = 0; c < num_channels; ++c) {
                float val = instance_seg_data[(y * width + x) * num_channels + c];
                normalized_sum += (val - min_vals[c]) / (max_vals[c] - min_vals[c]);
            }
            normalized_sum /= num_channels;
            instance_mask.at<uchar>(y, x) = static_cast<uchar>(normalized_sum * 255);
        }
    }

    // 确定转向方向
    int left_count = 0, right_count = 0;
    std::string turn_direction = determine_turn_direction(instance_seg_data, height, width, 
                                                        left_count, right_count);
    printf("L: %d, R: %d -> %s\n", left_count, right_count, turn_direction.c_str());

    // 创建可视化结果
    cv::Mat colored_mask;
    cv::applyColorMap(instance_mask, colored_mask, cv::COLORMAP_JET);
    cv::resize(colored_mask, colored_mask, original_image.size());

    // 创建叠加效果
    cv::Mat overlay;
    cv::addWeighted(original_image, 0.7, colored_mask, 0.3, 0, overlay);
    cv::imwrite("lane_detection_result.jpg", overlay);
}

int main() {
    try {
        // 加载优化后的模型
        std::unique_ptr<tflite::FlatBufferModel> model =
            tflite::FlatBufferModel::BuildFromFile("../src/lanenet_float32_optimized.tflite");
        if (!model) {
            throw std::runtime_error("Failed to load model");
        }
        printf("Model loaded successfully\n");

        // 创建解释器
        tflite::ops::builtin::BuiltinOpResolverWithXNNPACK resolver(true);
        std::unique_ptr<tflite::Interpreter> interpreter;
        
        if (tflite::InterpreterBuilder(*model, resolver)(&interpreter) != kTfLiteOk || !interpreter) {
            throw std::runtime_error("Failed to create interpreter");
        }

        // 设置线程数
        interpreter->SetNumThreads(4);

        // 设置输入张量大小
        interpreter->ResizeInputTensor(0, {1, 256, 512, 3});
        if (interpreter->AllocateTensors() != kTfLiteOk) {
            throw std::runtime_error("Failed to allocate tensors");
        }

        // 加载新的测试图片
        cv::Mat input_image = cv::imread("../src/2.png");
        if (input_image.empty()) {
            throw std::runtime_error("Failed to load image");
        }
        printf("Image loaded: %dx%d\n", input_image.cols, input_image.rows);

        // 预处理图片
        cv::Mat resized_image;
        //cv::resize(input_image, resized_image, cv::Size(512, 256));
        cv::Mat float_image;
        input_image.convertTo(float_image, CV_32F, 1.0/127.5, -1.0);
        
        // 复制数据到输入张量
        float* input_data = interpreter->typed_input_tensor<float>(0);
        memcpy(input_data, float_image.data, 256 * 512 * 3 * sizeof(float));

        // 执行推理
        printf("Running inference...\n");
        if (interpreter->Invoke() != kTfLiteOk) {
            throw std::runtime_error("Failed to invoke interpreter");
        }

        // 后处理输出
        printf("Processing outputs...\n");
        TfLiteTensor* instance_tensor = interpreter->tensor(interpreter->outputs()[1]);
        postprocess_outputs(instance_tensor->data.f, 256, 512, input_image);

        printf("Processing completed successfully\n");

    } catch (const std::exception& e) {
        printf("Error: %s\n", e.what());
        return 1;
    }

    return 0;
}